from connection import DBOptions,DBPool,db_pool
import unittest,json
from ddt import ddt,data
from loguru import logger
from datetime import date,datetime
import different
import settings

class JsonCustomEncoder(json.JSONEncoder):

    def default(self, value):
        if isinstance(value, datetime):
            return value.strftime('%Y-%m-%d %H:%M:%S')
        elif isinstance(value, date):
            return value.strftime('%Y-%m-%d')
        else:
            return json.JSONEncoder.default(self, value)

@ddt
class Checker(unittest.TestCase):
    original = settings.original
    ectypal = settings.ectypal
    # 忽略校验的表字段，eg：{表01：[字段01，字段02]}
    ignore = settings.ignore
    # 排序方式，默认排序方式在default中，指定排序方式在 custom 填写。 eg: {custom: { 表01: 排序字段，表02: 排序字段 }}
    order = settings.order
    where = settings.where

    def setUp(self) -> None:
        logger.add('run.log', )
        self.original_conn = db_pool(options=self.original)
        self.ectypal_conn = db_pool(options=self.ectypal)

    def tearDown(self) -> None:
        if self.original_conn: self.original_conn.close()
        if self.ectypal_conn: self.ectypal_conn.close()

    def test_library(self):
        original_data = [obj['Tables_in_' + self.original.db] \
         for obj in self.original_conn.query('show tables;', ()) \
         if obj['Tables_in_' + self.original.db] not in self.ignore['table']]

        ectypal_data = [obj['Tables_in_' + self.ectypal.db] \
         for obj in self.ectypal_conn.query('show tables;', ()) \
         if obj['Tables_in_' + self.ectypal.db] not in self.ignore['table']]

        set_msg = '\n==================== 原始库表数据 ====================\n'
        set_msg += json.dumps(original_data, cls=JsonCustomEncoder)
        set_msg += '\n==================== 迁移库表数据 ====================\n'
        set_msg += json.dumps(ectypal_data, cls=JsonCustomEncoder)
        set_msg += '\n===================================================='
        logger.debug(set_msg)
        original_table = set(original_data)
        ectypal_table = set(ectypal_data)
        intersection = list(original_table & ectypal_table)
        logger.debug('交集：'+ str(intersection) )
        no_migration = original_table - ectypal_table
        new_table = ectypal_table - original_table

        content_diff = {}
        for table in intersection:
            logger.debug('开始校验表格：{}', table)
            diff = self.table_check(table)
            if diff:
                content_diff[table] = diff

        if no_migration or new_table  or content_diff:
            error_msg = '\n==================== 未迁移表 ====================\n'
            error_msg += json.dumps(list(no_migration) or None, ensure_ascii=False, cls=JsonCustomEncoder )
            error_msg += '\n==================== 新增表 ====================\n'
            error_msg += json.dumps(list(new_table) or None, ensure_ascii=False, cls=JsonCustomEncoder)
            error_msg += '\n==================== 内容对比 ====================\n'
            error_msg += json.dumps(content_diff or None, ensure_ascii=False, cls=JsonCustomEncoder)
            error_msg += '\n===================================================='
            logger.error(error_msg)
        else:
            logger.success('原始库与迁移库数据一致')
            assert True

    def subset(self,sql:str, slice_):
        original = self.original_conn.query(sql, ())
        ectypal = self.ectypal_conn.query(sql, ())
        try:
            return different.batch_diff(original, ectypal)
        except different.LengthException as e:
            return ['在区间{}之间，对比数据长度不一致！'.format(str(slice_))]

    def table_check(self, table: str):
        sql = '''select %s from %s;'''
        order = self.order['custom'].get(table, self.order['default'])
        ignore = self.ignore['field'].get(table, [])
        where = self.where['custom'].get(table, self.where['default'])

        # 查询所以字段
        columns_sql = '''show full columns from %s;''' % (table,)
        logger.debug('获取字段名称 >>>')
        fields_data = self.original_conn.query(columns_sql, ())
        # 获取字段列表
        fields = [field['Field'] for field in fields_data]

        where_field = [key_value.split('=')[0] for key_value in where.upper().split('AND')]
        # 判断查询条件是否合法
        query = [table]
        for field in fields:
            if field.upper() in where_field:
                query.append('where')
                query.append(where)
            if field.upper() == order.upper():
                query.append('order by')
                query.append(order)

        if ignore:
            fields = list(set(fields) - set(ignore))

        total_sql = sql % ('count(1) as total', ' '.join(query))
        total_data = self.original_conn.query(total_sql, ())
        total = 0
        if total_data:
            total = total_data[0]['total']

        if total == 0:
            return []

        group = different.complete_slice(total)
        result = []
        diff_detail = []
        from concurrent.futures import ThreadPoolExecutor, as_completed
        threadPool = ThreadPoolExecutor(max_workers=10, thread_name_prefix="test_")
        for i in range(len(group)):
            data_sql = sql % (','.join(fields), ' '.join([*query, 'limit {}'.format(','.join([ str(i) for i in group[i]]))]))
            future = threadPool.submit(self.subset, data_sql, group[i])
            result.append(future)
        threadPool.shutdown(wait=True)
        for future in as_completed(result):
            data = future.result()
            if len(diff_detail) > settings.diff_max_count:
                return diff_detail
            diff_detail.extend(data)

        return diff_detail


if __name__ == '__main__':
    suite = unittest.TestLoader().loadTestsFromTestCase(Checker)
    unittest.TextTestRunner(verbosity=2).run(suite)